#!/bin/python3

import matplotlib.pylab as plt
import scipy.sparse as sparse
import random
import numpy as np
import time
from multiprocessing import Process

#TODO: Need to change.

def my_pic_para():
    my_pic = {}
    my_pic["style"]         = "png"
    my_pic["dpi"]           = 300
    my_pic["myfontsize"]    = 12
    my_pic["titlefontsize"] = 18
    my_pic["figsize"] = [4,3]
    my_pic["markersize"] = 0.5e-5
    my_pic["transparent"] = False

    # Solver01
    my_pic["save_pic_name"] = "solver01"
    my_pic["save_pic_dir"]  = "/home/matrix/solverchanllenge2021/solverchallenge21_01"
    my_pic["matrix_path"]   = "/home/matrix/solverchanllenge2021/solverchallenge21_01/jxpamg_matrix01"
    my_pic["save_matrix_path"]   = "/home/matrix/solverchanllenge2021/solverchallenge21_01/jxpamg_matrix01_change"
    
    # my_pic["save_pic_name"] = "shangfei"
    # my_pic["save_pic_dir"]  = "/home/matrix/iter00001/A"
    # my_pic["matrix_path"]   = "/home/matrix/iter00001/A/jxpamg_iter00001_job0"
    # my_pic["save_matrix_path"]   = "/home/matrix/iter00001/A/jxpamg_iter00001_job0_change"

    # my_pic["save_pic_name"] = "lixue"
    # my_pic["save_pic_dir"]  = "/home/matrix/hsl_matrix"
    # my_pic["matrix_path"]   = "/home/matrix/hsl_matrix/jxpamg_matrix"
    # my_pic["save_matrix_path"]   = "/home/matrix/hsl_matrix/lixue_change"

    # for test
    # my_pic["save_pic_name"] = "test"
    # my_pic["save_pic_dir"]  = "/home/mrz/code/python_script/pic_plot/plot_matrix"
    # my_pic["matrix_path"]   = "/home/mrz/code/python_script/pic_plot/plot_matrix/mat_20X20"
    # my_pic["save_matrix_path"]   = "/home/mrz/code/python_script/pic_plot/plot_matrix/mat_20X20_test"
    
    my_pic["font"] = {'family': 'Times New Roman', 'size': '{fs}'.format(fs=my_pic["titlefontsize"])}
    return my_pic

def read_data(my_pic):
    rptr = []
    col = []
    value = []
    i     = 1
    n     = 0
    nnz   = 0
    row_first = 0
    col_first = 0
    print("read matrix: \n" +my_pic["matrix_path"])
    for line in open(my_pic["matrix_path"], "r" ):
        line = line.strip("\n")
        if(i == 1):
            n = int(line)
        elif(i > 1 and i <= n + 2):
            if(i == 2):
                row_first = int(line)
            elif(i == n + 2):
                nnz = int(line) - row_first
            rptr.append(int(line) - row_first)
        elif(i > n + 2 and i < n + 3 + nnz):
            if(i == n + 3):
                col_first =int(line)
            col.append(int(line) - col_first)
        else:
            value.append(float(line))
            # print("%.15e" %(float(line)))
        i = i + 1
    print("read matrix success: n is {}, nnz is {}".format(n,nnz))

    my_pic["csr"] = {}
    my_pic["csr"]["n"] = n
    my_pic["csr"]["nnz"] = nnz
    my_pic["csr"]["rptr"] = rptr
    my_pic["csr"]["col"]  = col
    my_pic["csr"]["value"]= value

    return my_pic

def rearrange_matrix(my_pic):
    # Note:
    # 每行非零元数和元素值没有变动
    # 变化的只有列标
    n    = my_pic["csr"]["n"] 
    # nnz  = my_pic["csr"]["n"] 
    rptr = my_pic["csr"]["rptr"]
    # col  = my_pic["csr"]["col"] 
    my_order_tag = 0
    rearrange_col = []
    k = 0
    #numpy
    
    prepare_list = np.arange(n)
    for i in range(n):
        # T1 = time.perf_counter()
        for j in range(rptr[i], rptr[i+1]):
            if(i == j):
                my_order_tag = 1
                val_index =  j - rptr[i]
        if(my_order_tag):
            # T1 = time.perf_counter()
            prepare_list = np.delete(prepare_list, i)
            my_random = np.random.choice(prepare_list,rptr[i+1]-rptr[i] - 1)
            my_random = np.insert(my_random,val_index,i)
            prepare_list = np.insert(prepare_list,i,i)
            # T2 = time.perf_counter()
            # print('random和insert时间:%s ms' % ((T2 - T1)*1000))
        else:
            # T1 = time.perf_counter()
            my_random = np.random.choice(prepare_list,rptr[i+1]-rptr[i])
            # T2 = time.perf_counter()
            # print('random时间:%s ms' % ((T2 - T1)*1000))
        my_random = my_random.tolist()
        rearrange_col = rearrange_col + my_random
        if i == k*100000:
            print("random row " + str(i) + " success! ", end = "\r")
            k = k + 1
        # print(my_random)
        # print(col[rptr[i]:rptr[i+1]])
        # T2 = time.perf_counter()
        # print('a row time:%s ms' % ((T2 - T1)*1000))
    # for i in range(n):
    #     prepare_list = list(range(0,n))
    #     del prepare_list[i]
    #     for j in range(rptr[i], rptr[i+1]):
    #         if(i == j):
    #             my_order_tag = 1
    #     if(my_order_tag):
    #         my_random = random.sample(prepare_list,rptr[i+1]-rptr[i] - 1)
    #         my_random.insert(0,i)
    #     else:
    #         my_random = random.sample(prepare_list,rptr[i+1]-rptr[i])
    #     rearrange_col = rearrange_col + my_random
    #     if i == k*100000:
    #         print("random row " + str(i) + " success! ", end = "\r")
    #         k = k + 1
    print("success random! ")
    # print(rearrange_col)
    my_pic["csr"]["rearrange_col"] =  rearrange_col
    print("rerrange matrix success")
    return my_pic

def save_rearrange_matrix(my_pic):
    n    = my_pic["csr"]["n"]
    nnz  = my_pic["csr"]["nnz"]
    rptr = my_pic["csr"]["rptr"]
    col  = my_pic["csr"]["col"]
    value  = my_pic["csr"]["value"]

    with open(my_pic["save_matrix_path"], "w+" ) as fo:
        fo.write(str(n)+"\n")
        for i in range(0, n + 1):
            # fo.write(str(rptr[i])+"\n")
            fo.write("%d" %(rptr[i])+"\n")
        for i in range(0, nnz):
            # fo.write(str(col[i])+"\n")
            fo.write("%d" %(col[i])+"\n")
        for i in range(0, nnz):
            # fo.write(str(value[i])+"\n")
            fo.write("%.15e" %(value[i])+"\n")

def plot_my_pic(my_pic):
    # arr = np.array([my_pic["csr"]["rptr"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["value"]],
    # dtype=object)
    mat = sparse.csr_matrix((my_pic["csr"]["value"], my_pic["csr"]["col"], my_pic["csr"]["rptr"],), shape=(my_pic["csr"]["n"], my_pic["csr"]["n"]))
    plt.rcParams['font.sans-serif']=[my_pic["font"]["family"]]
    plt.figure(figsize=my_pic["figsize"],dpi=my_pic["dpi"])
    plt.spy(mat, marker=',', markersize = my_pic["markersize"])
    plt.savefig(my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"_raw."+my_pic["style"],bbox_inches='tight')
    print("save pic in " + my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"_raw."+my_pic["style"])
    plt.close()

def plot_rearrange_pic(my_pic):
    # arr = np.array([my_pic["csr"]["rptr"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["value"]],
    # dtype=object)
    mat = sparse.csr_matrix((my_pic["csr"]["value"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["rptr"],), shape=(my_pic["csr"]["n"], my_pic["csr"]["n"]))
    plt.rcParams['font.sans-serif']=[my_pic["font"]["family"]]
    plt.figure(figsize=my_pic["figsize"], dpi=my_pic["dpi"])
    plt.spy(mat, markersize = my_pic["markersize"])
    plt.savefig(my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"."+my_pic["style"],bbox_inches='tight')
    print("save pic in " + my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"."+my_pic["style"])
    plt.close()

def matrix_bandwith(mat):
    bandwidth = 0
    for i in np.arange(mat.nnz):
        bandwidth = max(bandwidth, abs(mat.row[i] - mat.col[i]))
    return bandwidth

# def matrix_bandwith_definition(mat):
#     bandwidth = 0
#     for i in np.arange(mat.nnz):
#         bandwidth = max(bandwidth, abs(mat.row[i] - mat.col[i]))
#     return bandwidth

def get_mat_feature(mat):
    #matrix_bandwidth = bandwidth(mat.todense())
    bandwidth = matrix_bandwith(mat)
    sparsity = float(mat.nnz) / float(mat.shape[0]*mat.shape[1])
    nnz_per_row = mat.nnz / mat.shape[0]
    print("=====================")
    print(">> matrix feature is:")
    print(">> bandwith upper is {0}".format(str(bandwidth)))
    print(">> sparsity is {0:.3e}".format(sparsity))
    print(">> nnz per row is {0}".format(str(nnz_per_row)))
    print("=====================")



def trans2coo(my_pic):
    mat = sparse.csr_matrix((my_pic["csr"]["value"], my_pic["csr"]["col"], my_pic["csr"]["rptr"],), shape=(my_pic["csr"]["n"], my_pic["csr"]["n"]))
    # my_pic["csr_mat"] = mat
    my_pic["coo_mat"] = mat.tocoo()
    return my_pic

#TODO:画图代码不够美观
#COPY石油大学的部分代码
def plot_heatmap(my_pic):
    # arr = np.array([my_pic["csr"]["rptr"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["value"]],
    # dtype=object)
    mat = my_pic["coo_mat"]
    plt.rcParams['font.sans-serif']=[my_pic["font"]["family"]]
    fig = plt.figure(figsize=my_pic["figsize"], dpi = my_pic["dpi"])
    ax = plt.gca()      #获取到当前坐标轴信息
    ax.xaxis.set_ticks_position('top')   #将X坐标轴移到上面
    ax.invert_yaxis()   #反转Y坐标轴
    plt.scatter(mat.row, mat.col,c = mat.data, s = my_pic["markersize"])
    plt.colorbar()
    plt.xlim(0,mat.shape[1])
    plt.ylim(mat.shape[0],0) 
    fig.savefig("{0}/{1}.{2}".format(my_pic["save_pic_dir"],my_pic["save_pic_name"],my_pic["style"]),transparent=my_pic["transparent"],bbox_inches='tight')
    plt.close(fig)
    print("save pic in " + my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"."+my_pic["style"])

def main():
    my_pic = my_pic_para()
    my_pic = read_data(my_pic)
    my_pic = trans2coo(my_pic)
    plot_heatmap(my_pic)
    get_mat_feature(my_pic["coo_mat"])
    #plot_my_pic(my_pic)

    # rerange matrix
    # my_pic = rearrange_matrix(my_pic)
    # save_rearrange_matrix(my_pic)
    # plot_rearrange_pic(my_pic)

def test():
    return

if __name__=="__main__":
    test()
